package com.stmt.analytic.graph.algorithm.lv;

import com.stmt.analytic.graph.algorithm.IdMapping;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

public class LouvainCommunities {

    private int[] global_cluster;

    private IdMapping idMapping;

    private Map<Integer, LouvainCommunity> communities;

    public LouvainCommunities(int[] global_cluster, IdMapping idMapping) {
        this.global_cluster = global_cluster;
        this.idMapping = idMapping;
        communities = new HashMap<>();
        for (int i = 0; i < global_cluster.length; i++) {
            int no = global_cluster[i];
            if (!communities.containsKey(no)) {
                communities.put(no, new LouvainCommunity(no));
            }
            LouvainCommunity louvainCommunity = communities.get(no);
            louvainCommunity.add(idMapping.toOriginalId(i));
        }
    }

    public Map<Integer, LouvainCommunity> getAll() {
        return communities;
    }

    @Override
    public String toString() {
        return "LouvainCommunities{" +
                "communities=" + communities +
                '}';
    }

    public static final class LouvainCommunity {
        private int no;
        private Collection<String> members;

        public LouvainCommunity(int no) {
            this.no = no;
            members = new ArrayList<>();
        }


        public Collection<String> getMembers() {
            return members;
        }

        public void add(String member) {
            members.add(member);
        }

        @Override
        public String toString() {
            return members.toString();
        }
    }
}
